import os
import pandas as pd
import glob
from datetime import datetime

def merge_qa_files(csvqaa_dir: str, csvqatotal_dir: str, output_dir: str):
    """合并同一水库的csvQA_generator.py和csvQAtotal_generator.py生成的问答对"""
    print("=== CSV问答对合并工具 ===")
    print(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    # 查找所有带'a'后缀的文件（csvQA_generator.py生成的）
    a_pattern = os.path.join(csvqaa_dir, "*a.xlsx")
    a_files = glob.glob(a_pattern)
    
    # 查找所有带'_total'后缀的文件（csvQAtotal_generator.py生成的）
    total_pattern = os.path.join(csvqatotal_dir, "*_total.xlsx")
    total_files = glob.glob(total_pattern)
    
    if not a_files and not total_files:
        print(f"在目录 {csvqaa_dir} 和 {csvqatotal_dir} 中未找到任何问答对文件")
        return
    
    print(f"找到 {len(a_files)} 个基础问答对文件")
    print(f"找到 {len(total_files)} 个统计问答对文件")
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建文件映射：基础名称 -> [a文件, total文件]
    file_mapping = {}
    
    # 处理带'a'后缀的文件
    for a_file in a_files:
        base_name = os.path.basename(a_file).replace('a.xlsx', '')
        if base_name not in file_mapping:
            file_mapping[base_name] = {'a_file': None, 'total_file': None}
        file_mapping[base_name]['a_file'] = a_file
    
    # 处理带'_total'后缀的文件
    for total_file in total_files:
        base_name = os.path.basename(total_file).replace('_total.xlsx', '')
        if base_name not in file_mapping:
            file_mapping[base_name] = {'a_file': None, 'total_file': None}
        file_mapping[base_name]['total_file'] = total_file
    
    total_merged_files = 0
    total_qa_pairs = 0
    
    # 合并每个水库的问答对
    for base_name, files in file_mapping.items():
        print(f"\n正在合并: {base_name}")
        
        merged_data = []
        current_index = 1
        
        # 读取基础问答对文件
        if files['a_file'] and os.path.exists(files['a_file']):
            try:
                df_a = pd.read_excel(files['a_file'])
                print(f"  读取基础问答对: {len(df_a)} 条")
                
                # 重新编号并添加到合并数据
                for _, row in df_a.iterrows():
                    row_dict = row.to_dict()
                    row_dict['标号'] = current_index
                    merged_data.append(row_dict)
                    current_index += 1
                    
            except Exception as e:
                print(f"  读取文件 {files['a_file']} 时出错: {e}")
        
        # 读取统计问答对文件
        if files['total_file'] and os.path.exists(files['total_file']):
            try:
                df_total = pd.read_excel(files['total_file'])
                print(f"  读取统计问答对: {len(df_total)} 条")
                
                # 重新编号并添加到合并数据
                for _, row in df_total.iterrows():
                    row_dict = row.to_dict()
                    row_dict['标号'] = current_index
                    merged_data.append(row_dict)
                    current_index += 1
                    
            except Exception as e:
                print(f"  读取文件 {files['total_file']} 时出错: {e}")
        
        # 保存合并后的文件
        if merged_data:
            output_file = os.path.join(output_dir, f"{base_name}.xlsx")
            
            # 如果输出文件已存在，先删除
            if os.path.exists(output_file):
                os.remove(output_file)
            
            try:
                # 创建DataFrame并保存
                df_merged = pd.DataFrame(merged_data)
                df_merged.to_excel(output_file, index=False)
                
                total_merged_files += 1
                file_qa_count = len(merged_data)
                total_qa_pairs += file_qa_count
                
                print(f"  合并完成: {base_name}.xlsx，包含 {file_qa_count} 个问答对")
                
            except Exception as e:
                print(f"  保存合并文件时出错: {e}")
        else:
            print(f"  {base_name} 没有找到有效的问答对数据")
    
    print(f"\n合并完成！")
    print(f"总共合并了 {total_merged_files} 个文件")
    print(f"总共包含 {total_qa_pairs} 个问答对")
    print(f"结果已保存到目录: {output_dir}")
    print(f"结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

def main():
    """主函数 - 合并CSV问答对文件"""
    print("=== CSV问答对合并工具 ===")
    
    # 设置输入目录和输出目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    csvqaa_dir = os.path.join(current_dir, "csvQAa")
    csvqatotal_dir = os.path.join(current_dir, "csvQAtotal")
    output_dir = os.path.join(current_dir, "csvQA")
    
    print(f"当前工作目录: {current_dir}")
    print(f"基础问答对输入目录: {csvqaa_dir}")
    print(f"统计问答对输入目录: {csvqatotal_dir}")
    print(f"输出目录: {output_dir}")
    print()
    
    # 检查输入目录是否存在
    if not os.path.exists(csvqaa_dir) and not os.path.exists(csvqatotal_dir):
        print(f"错误：输入目录都不存在 - {csvqaa_dir} 和 {csvqatotal_dir}")
        return
    
    try:
        # 合并问答对文件
        merge_qa_files(csvqaa_dir, csvqatotal_dir, output_dir)
        
    except Exception as e:
        print(f"合并过程中发生错误: {e}")

if __name__ == "__main__":
    main()